import java.util.PriorityQueue;
import java.util.Queue;

/**
 * 295. 数据流的中位数
 * https://leetcode-cn.com/problems/find-median-from-data-stream/
 */
public class Solutions_295 {
    public static void main(String[] args) {
        MedianFinder obj = new MedianFinder();
        obj.addNum(1);
        obj.addNum(2);
        double res1 = obj.findMedian();  // output: 1.5  ((1 + 2) / 2)
        System.out.println(res1);
        obj.addNum(3);
        double res2 = obj.findMedian();  // output: 2
        System.out.println(res2);
    }
}

class MedianFinder {
    private Queue<Integer> minHeap = null;
    private Queue<Integer> maxHeap = null;

    /** initialize your data structure here. */
    public MedianFinder() {
        // 小顶堆：堆顶元素是堆中最小的元素，出队时，堆顶元素出队
        minHeap = new PriorityQueue<>();
        // 大顶堆：堆顶元素是堆中最大的元素，出队时，堆顶元素出队
        maxHeap = new PriorityQueue<>((a, b) -> b - a);
    }

    /** 从数据流中添加一个整数到数据结构中。 */
    public void addNum(int num) {
        /**
         * 两个堆，互相交换元素，保证中位数一定是最大堆或最小堆的堆顶元素，或是两个堆的堆顶元素的平均数
         *     minHeap     maxHeap
         *       4            3
         *       5            2
         *                    1
         * 两个堆中的元素存储趋势：minHeap 中的元素，必然大于等于 maxHeap 中的元素，且 minHeap 中的元素是从小到大，
         * 而 maxHeap 中的元素是从大到小，那么中位数，只可能在堆顶元素中产生
         */
        if (minHeap.size() >= maxHeap.size()) {
            // 添加元素到大顶堆（添加的元素不一定是 num，而是先将 num 添加到 minHeap 中，再取出 minHeap 堆顶元素，即最小元素）
            minHeap.offer(num);
            maxHeap.offer(minHeap.poll());
        } else {
            // 添加元素到小顶堆（添加的元素不一定是 num，而是先将 num 添加到 maxHeap 中，再取出 maxHeap 堆顶元素，即最大元素）
            maxHeap.offer(num);
            minHeap.offer(maxHeap.poll());
        }
    }

    /** 返回目前所有元素的中位数。 */
    public double findMedian() {
        if (minHeap.size() > maxHeap.size()) {
            // 中位数在小顶堆的堆顶
            return minHeap.peek();
        } else if (maxHeap.size() > minHeap.size()) {
            // 中位数在大顶堆的堆顶
            return maxHeap.peek();
        } else {
            // 两个中间元素的平均值
            return (minHeap.peek() + maxHeap.peek()) / 2.0;
        }
    }
}
